import logging
import json
from pathlib import Path
from typing import Any, Type

from pydantic import BaseModel, TypeAdapter

from .data import (Turn, Session, QA)


logger = logging.getLogger(__name__)


def _save_data(data: list[Type[BaseModel]], t: Any, file_path: str, filename: str):
    adapter = TypeAdapter(t)
    json_bytes = adapter.dump_json(data)
    complete_file_path = Path(file_path) / filename
    with complete_file_path.open("wb") as f:
        f.write(json_bytes)


def _load_data(t: Any, file_path: str, filename: str):
    adapter = TypeAdapter(t)
    complete_file_path = Path(file_path) / filename
    with complete_file_path.open("rb") as f:
        data = adapter.validate_python(json.load(f))
    return data


class LOCOMO:
    def __init__(self, data_dir: str, reload: bool = False):
        self.data_dir = data_dir
        self.conversations: list[list[Session]] = []
        self.qas: list[list[QA]] = []
        self.is_existed = Path(self.data_dir + "/conversations.json").exists() \
            and Path(self.data_dir + "/qas.json").exists()
        if not self.is_existed or reload:
            logger.info("processing locomo dataset")
            data_path = Path(self.data_dir) / "locomo10.json"
            self.datas = json.load(Path(data_path).open(mode="r", encoding="utf-8"))
            self._process_data()

        if self.is_existed and not reload:
            logger.info("loading locomo dataset from json file")
            self.conversations = _load_data(list[list[Session]], self.data_dir, "conversations.json")
            self.qas = _load_data(list[list[QA]], self.data_dir, "qas.json")

    def _process_data(self):
        for data in self.datas:
            i = 1
            qas, conversation = data["qa"], data["conversation"]
            # process the conversation
            sessions = []
            while session := conversation.get(f"session_{i}", []):
                turns = []
                for j in range(len(session)):
                    if blip_caption := session[j].get("blip_caption", ""):
                        turns.append(Turn(
                            speaker=session[j].get("speaker"),
                            dia_id=session[j].get("dia_id"),
                            text=session[j].get("text"),
                            blip_caption=blip_caption
                        ))
                    else:
                        turns.append(Turn(
                            speaker=session[j].get("speaker"),
                            dia_id=session[j].get("dia_id"),
                            text=session[j].get("text")
                        ))
                sessions.append(Session(
                    turns=turns,
                    datetime=conversation.get(f"session_{i}_date_time")
                ))
                i += 1
            self.conversations.append(sessions)
            # process the qa
            qas_list = []
            for qa in qas:
                if int(qa.get("category")) == 5:
                    adversarial_answer = qa.get("adversarial_answer")
                    qas_list.append(QA(
                        question=qa.get("question"),
                        evidence=qa.get("evidence"),
                        category=qa.get("category"),
                        adversarial_answer=adversarial_answer
                    ))
                else:
                    answer = qa.get("answer", "")
                    qas_list.append(QA(
                        question=qa.get("question"),
                        answer=str(answer),
                        evidence=qa.get("evidence"),
                        category=qa.get("category")
                    ))

            self.qas.append(qas_list)
        logger.info("save locomo dataset to json file")
        _save_data(self.conversations, list[list[Session]], self.data_dir, "conversations.json")
        _save_data(self.qas, list[list[QA]], self.data_dir, "qas.json")

    def get_evidence(self, id: str, index: int):
        conversation = self.conversations[index]
        session_id, dia_id = int(id.split(":")[0][-1]) - 1, int(id.split(":")[1]) - 1
        session = conversation[session_id]
        turn = session.turns[dia_id]
        return turn.text
